import os
import subprocess
from PIL import Image

# 定义需要遍历的图片文件扩展名
image_extensions = ['.png', '.jpg', '.jpeg', '.webp']

# 输出文件路径
output_file = 'image_list.txt'

def compress_png(image_path):
    try:
        # 使用pngquant压缩PNG图片
        subprocess.run(['pngquant', '--force', '--ext', '.png', image_path])
        return True
    except Exception as e:
        print(f'在压缩图片 {image_path} 时发生异常:', str(e))
        return False

def compress_webp(image_path):
    try:
        # 使用cwebp压缩WebP图片
        subprocess.run(['cwebp', '-q', '80', image_path, '-o', image_path])
        return True
    except Exception as e:
        print(f'在压缩图片 {image_path} 时发生异常:', str(e))
        return False

def compress_jpeg(image_path):
    try:
        # 打开JPEG图片
        image = Image.open(image_path)

        # 压缩JPEG图片
        image.save(image_path, optimize=True, quality=80)

        # 关闭图片
        image.close()

        return True
    except Exception as e:
        print(f'在压缩图片 {image_path} 时发生异常:', str(e))
        return False

def modify_exif(image_path):
    try:
        # 打开图片
        image = Image.open(image_path)

        # 将参数值写入图片的 EXIF 数据中
        exif_data = image.info.get('exif', b'')
        exif_data += b'SpecialParameter=party_go\0'
        image.info['exif'] = exif_data

        # 保存修改后的图片
        image.save(image_path)

        # 关闭图片
        image.close()

        return True
    except Exception as e:
        print(f'在修改图片 {image_path} 的EXIF数据时发生异常:', str(e))
        return False

def append_info_to_file_footer(file_path, string):
    try:
        with open(file_path, 'a') as f:
            f.write(str(string))
        return True
    except Exception as e:
        print(f'在写入文件 {file_path} 时发生异常:', str(e))
        return False


def process_images():
    # 遍历当前目录及其子目录下的所有文件
    with open(output_file, 'w') as f:
        for root, dirs, files in os.walk('.'):
            for file in files:
                # 获取文件的扩展名
                ext = os.path.splitext(file)[1].lower()
                # 检查扩展名是否为图片文件
                if ext in image_extensions:
                    # 图片路径
                    image_path = os.path.join(root, file)

                    # 根据文件扩展名选择压缩方法
                    if ext == '.png':
                        # 压缩PNG图片
                        if compress_png(image_path):
                            # 输出图片文件的路径及所在目录到文本文件
                            f.write(f'{image_path}\t{root}\n')
                            append_info_to_file_footer(image_path,"Power By Party Go")
                    elif ext == '.webp':
                        # 压缩WebP图片
                        if compress_webp(image_path):
                            # 输出图片文件的路径及所在目录到文本文件
                            f.write(f'{image_path}\t{root}\n')
                            append_info_to_file_footer(image_path, "Power By Party Go")
                    elif ext in ['.jpg', '.jpeg']:
                        # 压缩JPEG图片
                        if compress_jpeg(image_path):
                            # 输出图片文件的路径及所在目录到文本文件
                            f.write(f'{image_path}\t{root}\n')
                            append_info_to_file_footer(image_path, "Power By Party Go")

    print('图片已压缩并修改EXIF数据')
    print('图片列表已保存到', output_file)

if __name__ == "__main__":
    process_images()